## 测试文件，数据已剔除
import csv
import gensim
from gensim.models import Word2Vec
from gensim.models import word2vec
from gensim.test.utils import common_texts
from nltk.util import pr


def get_word_feature():
    with open('data/table_names/newsapp_tablename_old.csv', encoding='utf-8-sig') as csvfile:
        reader = csv.reader(csvfile)
        text = []
        for row in reader:
            text.append(row[2].split('_'))

    # print(text)
    # print(common_texts)

    csvfile.close()

    # print(text)
    def rule(word, count, min_count):  # 设置哪些词排除在外
        if word in ['a', 'b', 'w', 'dw']:
            return gensim.utils.RULE_DEFAULT
        else:
            return gensim.utils.RULE_KEEP

    model = Word2Vec(text, size=250, window=1, min_count=1, trim_rule=rule)
    # Word2Vec(sentences=None,size=100,alpha=0.025,window=5, min_count=5, max_vocab_size=None, sample=0.001,seed=1, workers=3,min_alpha=0.0001, sg=0, hs=0, negative=5, cbow_mean=1, hashfxn=<built-in function hash>,iter=5,null_word=0, trim_rule=None, sorted_vocab=1, batch_words=10000)
    model.save('models/word_feature.model')
    return


def test_mdodel():
    model = Word2Vec.load('models/word_feature.model')
    print(model.similarity('sd', 'cgi'))


if __name__ == '__main__':
    get_word_feature()
    test_mdodel()
